package com.cloudsubassemblykafka10024.partitioner;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description: 用户自定义分区
 * @Author zx
 * @Date 2024/12/24 20:35
 * @Version 1.0
 */
public class UserDefinePartitioner implements Partitioner {

    // 模拟一个计数器
    private AtomicInteger counter = new AtomicInteger(0);

    /**
     * 返回分区号
     * @param topic The topic name
     * @param key The key to partition on (or null if no key)
     * @param keyBytes The serialized key to partition on( or null if no key)
     * @param value The value to partition on or null
     * @param valueBytes The serialized value to partition on or null
     * @param cluster The current cluster metadata
     * @return
     */
    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {

        // 获取分区信息
        List<PartitionInfo> partitions = cluster.availablePartitionsForTopic(topic);
        // 获取分区数量
        int size = partitions.size();
        if (keyBytes == null) {
            int andIncrement = counter.getAndIncrement();
            // 使用取模算法
            return (andIncrement & Integer.MAX_VALUE) % size;
        } else {
            // 使用hash算法
            return Utils.toPositive(Utils.murmur2(keyBytes)) % size;
        }

    }

    @Override
    public void close() {
        System.out.println("close");
    }

    @Override
    public void configure(Map<String, ?> configs) {
        System.out.println("configure");
    }
}
